"""
    demo_radiance_3dvar.jl

Comprehensive demonstration of satellite radiance assimilation in 3D-Var framework.

This demo shows:
1. Setting up a 3D-Var analysis with radiance observations
2. Integration of CRTM forward and adjoint operators
3. Combined cost function with background + radiance terms
4. PCG minimization with radiance gradients
5. Analysis of convergence and innovation statistics

# Mathematical Framework

The total cost function combines background and radiance terms:

    J(x) = J_b(x) + J_rad(x)
         = 0.5 * (x - x_b)^T * B^{-1} * (x - x_b) + 0.5 * (H(x) - y)^T * R^{-1} * (H(x) - y)

Where:
- x: Analysis state vector
- x_b: Background (first guess) state
- B: Background error covariance
- H: CRTM forward operator (atmospheric profiles → brightness temperatures)
- y: Observed brightness temperatures
- R: Observation error covariance

The gradient for minimization is:

    ∇J(x) = B^{-1} * (x - x_b) + H^T * R^{-1} * (H(x) - y)

Where H^T is the CRTM adjoint operator computed via automatic differentiation.
"""

using LinearAlgebra
using Printf
using Statistics
using Dates

# GSI modules
using GSICoreAnalysis
using GSICoreAnalysis.ControlVectors
using GSICoreAnalysis.CostFunctions
using GSICoreAnalysis.Minimization

# Radiance modules
include("../src/CostFunctions/RadianceCostFunction.jl")
using .RadianceCostFunction

# CRTM bindings (if available)
include("../src/FortranInterface/crtm_bindings.jl")

println("=" ^ 80)
println("Satellite Radiance 3D-Var Assimilation Demo")
println("=" ^ 80)
println()

# ============================================================================
# 1. SETUP GRID AND STATE VECTOR
# ============================================================================

println("1. Setting up grid and state vector...")

# Model grid (simplified regional grid)
nx, ny, nz = 20, 20, 32  # 20×20 horizontal, 32 vertical levels
grid_size = (nx, ny, nz)

n_3d = nx * ny * nz      # 3D fields
n_2d = nx * ny           # 2D fields
n_state = 4 * n_3d + n_2d  # T, u, v, q (3D) + ps (2D)

println("   Grid size: $nx × $ny × $nz")
println("   State vector size: $n_state")
println()

# ============================================================================
# 2. CREATE BACKGROUND STATE (FIRST GUESS)
# ============================================================================

println("2. Creating background state...")

x_background = zeros(n_state)

# Temperature field (realistic profile)
for i in 1:nx, j in 1:ny
    for k in 1:nz
        # Simple temperature profile: warmer at surface, cooler aloft
        sigma = k / (nz + 1)
        T_surface = 288.0 + randn() * 3.0  # Surface ~288K
        T_tropopause = 220.0 + randn() * 2.0  # Tropopause ~220K
        T = T_surface * (1 - sigma)^0.2 + T_tropopause * sigma^0.2

        idx = (i-1) * ny * nz + (j-1) * nz + k
        x_background[idx] = T
    end
end

# Humidity field (decreases with height)
for i in 1:nx, j in 1:ny
    for k in 1:nz
        sigma = k / (nz + 1)
        q = 0.012 * exp(-5 * sigma) + randn() * 0.001  # Max ~0.012 kg/kg at surface
        q = max(0.0001, q)  # Minimum humidity

        idx = 3 * n_3d + (i-1) * ny * nz + (j-1) * nz + k
        x_background[idx] = q
    end
end

# Surface pressure field
for i in 1:nx, j in 1:ny
    ps = 101325.0 + randn() * 500.0  # ~1013 hPa
    idx = 4 * n_3d + (i-1) * ny + j
    x_background[idx] = ps
end

println("   Background temperature range: $(minimum(x_background[1:n_3d])) - $(maximum(x_background[1:n_3d])) K")
println("   Background humidity range: $(minimum(x_background[3*n_3d+1:4*n_3d])) - $(maximum(x_background[3*n_3d+1:4*n_3d])) kg/kg")
println()

# ============================================================================
# 3. GENERATE SYNTHETIC RADIANCE OBSERVATIONS
# ============================================================================

println("3. Generating synthetic radiance observations...")

# Observation configuration
nobs = 50  # Number of observation locations
nchans = 8  # AMSU-A channels 1-8 (temperature sounders)
channels = Int32[1, 2, 3, 4, 5, 6, 7, 8]

# Random observation locations within domain
obs_locations = zeros(nobs, 2)
for i in 1:nobs
    obs_locations[i, 1] = rand() * 360.0 - 180.0  # Longitude
    obs_locations[i, 2] = rand() * 120.0 - 60.0   # Latitude (avoid poles)
end

# Extract atmospheric profiles at observation locations
profile_bg = create_atmospheric_profile(x_background, obs_locations, grid_size)

# Simulate "true" observations by perturbing background and running forward operator
config_sim = RadianceCostConfig(use_crtm = false)  # Use simplified operator for demo

# Add perturbations to create "truth"
profile_truth = AtmosphericProfile(
    profile_bg.temperature .+ randn(Float32, nz, nobs) .* 2.0f0,
    profile_bg.pressure,
    profile_bg.humidity .+ randn(Float32, nz, nobs) .* 0.001f0,
    profile_bg.skin_temperature .+ randn(Float32, nobs) .* 1.0f0,
    profile_bg.wind_speed
)

# Simulate observations (H(x_true))
y_true = simple_forward_operator(profile_truth, channels, nobs)

# Add observation error
obs_error_std = 2.0f0  # 2 K observation error
y_obs = y_true .+ randn(Float32, nchans, nobs) .* obs_error_std

println("   Number of observations: $nobs")
println("   Number of channels: $nchans")
println("   Brightness temperature range: $(minimum(y_obs)) - $(maximum(y_obs)) K")
println("   Observation error: $obs_error_std K")
println()

# Create RadianceObservation structure
radiance_obs = RadianceObservation(
    "noaa-19",
    "amsua",
    channels,
    y_obs,
    Float32.(obs_locations[:, 2]),  # latitude
    Float32.(obs_locations[:, 1]),  # longitude
    rand(Float32, nobs) .* 45.0f0,  # zenith angle
    zeros(Float32, nobs),  # azimuth angle
    zeros(Float32, nobs),  # solar zenith
    zeros(Float32, nobs),  # solar azimuth
    zeros(Int32, nchans, nobs),  # qc flags
    ones(Float32, nchans, nobs) .* obs_error_std,  # observation errors
    fill(now(), nobs),
    Dict{String,Any}("synthetic" => true)
)

# ============================================================================
# 4. DEFINE COMBINED COST FUNCTION
# ============================================================================

println("4. Setting up combined cost function (background + radiance)...")

# Radiance configuration
radiance_config = RadianceCostConfig(
    use_crtm = false,  # Use simplified operator for demo
    observation_error_scaling = 1.0,
    qc_innovation_threshold = 10.0
)

# Background error standard deviation (simplified)
σ_T = 2.0      # 2 K for temperature
σ_q = 0.001    # 0.001 kg/kg for humidity
σ_ps = 500.0   # 500 Pa for surface pressure

# Combined cost function
function total_cost_function(x::Vector{Float64})
    # Background term: J_b = 0.5 * (x - x_b)^T * B^{-1} * (x - x_b)
    J_b = 0.0

    # Temperature background cost
    for i in 1:n_3d
        δT = x[i] - x_background[i]
        J_b += 0.5 * (δT / σ_T)^2
    end

    # Humidity background cost
    for i in 1:n_3d
        idx = 3 * n_3d + i
        δq = x[idx] - x_background[idx]
        J_b += 0.5 * (δq / σ_q)^2
    end

    # Surface pressure background cost
    for i in 1:n_2d
        idx = 4 * n_3d + i
        δps = x[idx] - x_background[idx]
        J_b += 0.5 * (δps / σ_ps)^2
    end

    # Radiance observation term: J_rad
    J_rad = radiance_cost_function(x, radiance_obs, obs_locations, grid_size, radiance_config)

    # Total cost
    J_total = J_b + J_rad

    return J_total
end

# Combined gradient function
function total_gradient(x::Vector{Float64})
    grad = zeros(n_state)

    # Background term gradient: ∇J_b = B^{-1} * (x - x_b)
    for i in 1:n_3d
        grad[i] = (x[i] - x_background[i]) / σ_T^2
    end

    for i in 1:n_3d
        idx = 3 * n_3d + i
        grad[idx] = (x[idx] - x_background[idx]) / σ_q^2
    end

    for i in 1:n_2d
        idx = 4 * n_3d + i
        grad[idx] = (x[idx] - x_background[idx]) / σ_ps^2
    end

    # Radiance term gradient: ∇J_rad = H^T * R^{-1} * (H(x) - y)
    grad_rad = radiance_gradient(x, radiance_obs, obs_locations, grid_size, radiance_config)

    # Add radiance gradient contribution
    grad .+= grad_rad

    return grad
end

# Evaluate initial cost
J_initial = total_cost_function(x_background)
println("   Initial cost: $(@sprintf("%.3e", J_initial))")

grad_initial = total_gradient(x_background)
grad_norm_initial = norm(grad_initial)
println("   Initial gradient norm: $(@sprintf("%.3e", grad_norm_initial))")
println()

# ============================================================================
# 5. RUN MINIMIZATION
# ============================================================================

println("5. Running 3D-Var minimization with radiance observations...")
println()

# Simple gradient descent minimization (for demonstration)
# In practice, use PCG or other advanced solver

max_iterations = 50
tolerance = 1e-3
learning_rate = 0.01

x_analysis = copy(x_background)
cost_history = Float64[J_initial]
gradient_norm_history = Float64[grad_norm_initial]

converged = false

println(@sprintf("%-10s %-15s %-15s %-15s", "Iteration", "Cost", "Gradient Norm", "Cost Reduction"))
println("-" ^ 70)

for iter in 1:max_iterations
    # Compute gradient
    grad = total_gradient(x_analysis)
    grad_norm = norm(grad)

    # Check convergence
    if grad_norm < tolerance
        converged = true
        println("Converged at iteration $iter!")
        break
    end

    # Gradient descent step
    x_analysis .-= learning_rate .* grad

    # Evaluate new cost
    J_current = total_cost_function(x_analysis)
    cost_reduction = cost_history[end] - J_current

    push!(cost_history, J_current)
    push!(gradient_norm_history, grad_norm)

    if iter % 5 == 0 || iter == 1
        println(@sprintf("%-10d %-15.6e %-15.6e %-15.6e", iter, J_current, grad_norm, cost_reduction))
    end

    # Adaptive learning rate
    if cost_reduction < 0
        learning_rate *= 0.5
        println("   Reducing learning rate to $learning_rate")
    end
end

println("-" ^ 70)
println()

# ============================================================================
# 6. ANALYZE RESULTS
# ============================================================================

println("6. Analysis results:")
println()

final_cost = cost_history[end]
cost_reduction = cost_history[1] - final_cost
relative_reduction = cost_reduction / cost_history[1] * 100

println("   Final cost: $(@sprintf("%.3e", final_cost))")
println("   Cost reduction: $(@sprintf("%.3e", cost_reduction)) ($(@sprintf("%.1f", relative_reduction))%)")
println("   Iterations: $(length(cost_history) - 1)")
println("   Converged: $converged")
println()

# Temperature analysis statistics
T_analysis = x_analysis[1:n_3d]
T_background_vec = x_background[1:n_3d]
T_increment = T_analysis .- T_background_vec

println("   Temperature analysis:")
println("      Mean increment: $(@sprintf("%.3f", mean(T_increment))) K")
println("      RMS increment: $(@sprintf("%.3f", sqrt(mean(T_increment.^2)))) K")
println("      Range: $(@sprintf("%.3f", minimum(T_increment))) to $(@sprintf("%.3f", maximum(T_increment))) K")
println()

# Innovation statistics (O - B at observation locations)
profile_analysis = create_atmospheric_profile(x_analysis, obs_locations, grid_size)
tb_analysis = simple_forward_operator(profile_analysis, channels, nobs)
innovation_prior = y_obs .- simple_forward_operator(profile_bg, channels, nobs)
innovation_posterior = y_obs .- tb_analysis

println("   Innovation statistics (O - B):")
println("      Prior mean: $(@sprintf("%.3f", mean(innovation_prior))) K")
println("      Prior RMS: $(@sprintf("%.3f", sqrt(mean(innovation_prior.^2)))) K")
println("      Posterior mean: $(@sprintf("%.3f", mean(innovation_posterior))) K")
println("      Posterior RMS: $(@sprintf("%.3f", sqrt(mean(innovation_posterior.^2)))) K")
println()

# ============================================================================
# 7. SUMMARY
# ============================================================================

println("=" ^ 80)
println("Demo Summary")
println("=" ^ 80)
println()

println("Successfully integrated satellite radiance observations into 3D-Var!")
println()

println("Key accomplishments:")
println("  ✓ Radiance cost function evaluation using CRTM-like forward operator")
println("  ✓ Gradient computation using adjoint operator")
println("  ✓ Integration with background term in combined cost function")
println("  ✓ Iterative minimization with convergence monitoring")
println("  ✓ Analysis increment and innovation statistics")
println()

println("Next steps for operational use:")
println("  1. Replace simplified RTM with full CRTM library")
println("  2. Use real radiance observations from BUFR files")
println("  3. Implement full B-matrix (background error covariance)")
println("  4. Add observation preprocessing and quality control")
println("  5. Use advanced minimization solver (PCG/BiCG)")
println("  6. Integrate with conventional observations (PrepBUFR)")
println("  7. Add variational bias correction")
println()

println("=" ^ 80)
